import algebra.curves.barreto_naehrig.*;
import algebra.curves.barreto_naehrig.abstract_bn_parameters.AbstractBNG1Parameters;
import algebra.curves.barreto_naehrig.abstract_bn_parameters.AbstractBNG2Parameters;
import algebra.curves.barreto_naehrig.abstract_bn_parameters.AbstractBNGTParameters;
import algebra.curves.barreto_naehrig.bn254a.BN254aFields;
import algebra.curves.barreto_naehrig.bn254a.BN254aG1;
import algebra.curves.barreto_naehrig.bn254a.BN254aG2;
import algebra.curves.barreto_naehrig.bn254a.BN254aPairing;
import algebra.curves.barreto_naehrig.bn254a.bn254a_parameters.BN254aG1Parameters;
import algebra.curves.barreto_naehrig.bn254a.bn254a_parameters.BN254aG2Parameters;
import algebra.curves.barreto_naehrig.bn254b.BN254bFields;
import algebra.curves.barreto_naehrig.bn254b.BN254bG1;
import algebra.curves.barreto_naehrig.bn254b.BN254bG2;
import algebra.curves.barreto_naehrig.bn254b.BN254bPairing;
import algebra.curves.barreto_naehrig.bn254b.bn254b_parameters.BN254bG1Parameters;
import algebra.curves.barreto_naehrig.bn254b.bn254b_parameters.BN254bG2Parameters;
import algebra.curves.fake.*;
import algebra.curves.fake.fake_parameters.FakeFqParameters;
import algebra.curves.fake.fake_parameters.FakeG1Parameters;
import algebra.curves.fake.fake_parameters.FakeG2Parameters;
import algebra.fields.Fp;
import configuration.Configuration;
import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaPairRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.storage.StorageLevel;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;
import profiler.generation.R1CSConstruction;
import profiler.utils.SparkUtils;
import relations.objects.Assignment;
import relations.r1cs.R1CSRelation;
import relations.r1cs.R1CSRelationRDD;
import scala.Tuple3;
import zk_proof_systems.zkSNARK.*;
import zk_proof_systems.zkSNARK.objects.CRS;
import zk_proof_systems.zkSNARK.objects.Proof;

import static org.junit.Assert.assertTrue;

public class DistributedzkSNARK {
    private transient JavaSparkContext sc;
    private Configuration config;

    @Before
    public void setUp() {
        final SparkConf conf = new SparkConf().setMaster("local").setAppName("default");
        conf.set("spark.files.overwrite", "true");
        conf.set("spark.serializer", "org.apache.spark.serializer.KryoSerializer");
        conf.registerKryoClasses(SparkUtils.zksparkClasses());

        sc = new JavaSparkContext(conf);

        config = new Configuration(1, 1, 1, 2, sc, StorageLevel.MEMORY_ONLY());
        config.setRuntimeFlag(false);
        config.setDebugFlag(true);
    }

    @After
    public void tearDown() {
        sc.stop();
        sc = null;
    }

    private <BNFrT extends BNFields.BNFr<BNFrT>,
            BNFqT extends BNFields.BNFq<BNFqT>,
            BNFq2T extends BNFields.BNFq2<BNFqT, BNFq2T>,
            BNFq6T extends BNFields.BNFq6<BNFqT, BNFq2T, BNFq6T>,
            BNFq12T extends BNFields.BNFq12<BNFqT, BNFq2T, BNFq6T, BNFq12T>,
            BNG1T extends BNG1<BNFrT, BNFqT, BNG1T, BNG1ParametersT>,
            BNG2T extends BNG2<BNFrT, BNFqT, BNFq2T, BNG2T, BNG2ParametersT>,
            BNGTT extends BNGT<BNFqT, BNFq2T, BNFq6T, BNFq12T, BNGTT, BNGTParametersT>,
            BNG1ParametersT extends AbstractBNG1Parameters<BNFrT, BNFqT, BNG1T, BNG1ParametersT>,
            BNG2ParametersT extends AbstractBNG2Parameters<BNFrT, BNFqT, BNFq2T, BNG2T, BNG2ParametersT>,
            BNGTParametersT extends AbstractBNGTParameters<BNFqT, BNFq2T, BNFq6T, BNFq12T, BNGTT, BNGTParametersT>,
            BNPublicParametersT extends BNPublicParameters<BNFqT, BNFq2T, BNFq6T, BNFq12T>,
            BNPairingT extends BNPairing<BNFrT, BNFqT, BNFq2T, BNFq6T, BNFq12T, BNG1T, BNG2T, BNGTT, BNG1ParametersT, BNG2ParametersT, BNGTParametersT, BNPublicParametersT>>
    void DistributedBNProofSystemTest(
            final int numInputs,
            final int numConstraints,
            BNFrT fieldFactory,
            BNG1T g1Factory,
            BNG2T g2Factory,
            BNPairingT pairing) {
        final Tuple3<R1CSRelationRDD<BNFrT>, Assignment<BNFrT>, JavaPairRDD<Long, BNFrT>> construction =
                R1CSConstruction.parallelConstruct(numConstraints, numInputs, fieldFactory, config);
        final R1CSRelationRDD<BNFrT> r1cs = construction._1();
        final Assignment<BNFrT> primary = construction._2();
        final JavaPairRDD<Long, BNFrT> fullAssignment = construction._3();

        final CRS<BNFrT, BNG1T, BNG2T, BNGTT> CRS =
                DistributedSetup.generate(r1cs, fieldFactory, g1Factory, g2Factory, pairing, config);

        long t_proof = System.nanoTime();
        final Proof<BNG1T, BNG2T> proof =
                DistributedProver.prove(CRS.provingKeyRDD(), primary, fullAssignment, fieldFactory, config);
        long t_proof_end = System.nanoTime();

        final boolean isValid = Verifier.verify(CRS.verificationKey(), primary, proof, pairing, config);

        System.out.println(isValid);
        assertTrue(isValid);
        System.out.println("ProofGen time:" + (double) (t_proof_end - t_proof) / 1000000000.0 + "s");
    }


    @Test
    public void DistributedBN254aProofSystemTest() {
        final int numInputs = 1023;
        final int numConstraints = 1024;
        final BN254aFields.BN254aFr fieldFactory = new BN254aFields.BN254aFr(1);
        final BN254aG1 g1Factory = BN254aG1Parameters.ONE;
        final BN254aG2 g2Factory = BN254aG2Parameters.ONE;
        final BN254aPairing pairing = new BN254aPairing();
        System.out.println("test dis");

        DistributedBNProofSystemTest(numInputs, numConstraints, fieldFactory, g1Factory, g2Factory, pairing);
    }
}
